import logging
import os
import time

from _dynamic_dataset.AccCombinationRecord import WriteRecordsToFile


def Algorithm_MR_DD(num_epochs, clients, server, device, beta, sample_ratio_list, combinations, test_example_num, dirPath,
                    acc_record, time_record, ignore_rate_record, eval_type):
    os.makedirs(dirPath, exist_ok=True)
    logging.basicConfig(filename=f'{dirPath}/main.log', level=logging.INFO)
    # 开始时间
    start_time = time.time()
    for epoch in range(1, num_epochs + 1):
        logging.info(f"epoch={epoch}")
        # 客户端训练
        train_start_time = time.time()
        for client in clients:
            client.train(server.global_model)
        train_end_time = time.time()
        # 这里用(1,)这个组合来代表该轮的贡献度评估时间
        time_record.addCombinationAcc(epoch, (1,), train_end_time - train_start_time)
        test_start_time = time.time()
        # 服务器梯度聚合并计算对应的正确率
        for combination in combinations:
            # 创建客户端子集
            com_clients = []
            for clientId in combination:
                com_clients.append(clients[clientId])
            server.model_aggregate(com_clients, server.sub_model)
            # 在测试集上进行测试
            com_client_num = len(com_clients)
            fit_precision, ir = server.contribution_eval(server.sub_model, server.eval_loader, device, eval_type,
                                                         sample_ratio_list[com_client_num], com_client_num,
                                                         test_example_num,beta)
            acc_record.addCombinationAcc(epoch, combination, fit_precision)
            ignore_rate_record.addCombinationAcc(epoch, combination, ir)
        test_end_time = time.time()
        # 这里用(2,)这个组合来代表该轮的贡献度评估时间
        time_record.addCombinationAcc(epoch, (2,), test_end_time - test_start_time)
        # 更新全局模型、sub_model
        server.model_aggregate(clients, server.global_model)
    # 结束时间
    end_time = time.time()
    logging.info(f"总耗时为{end_time - start_time}s")
    WriteRecordsToFile(dirPath + "/acc_records.json", acc_record)
    WriteRecordsToFile(dirPath + "/ignore_rate_record.json", ignore_rate_record)
    WriteRecordsToFile(dirPath + "/time_records.json", time_record)
